package com.lee.windowFunction;

import com.lee.IP;
import com.lee.TCP;
import com.lee.UDP;
import com.lee.networkData;
import com.lee.windowFunction.StatisticalCountWindowFunction;
import org.apache.flink.streaming.api.functions.windowing.AllWindowFunction;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;

/**
 * Created by liruopeng on 2018/1/18.
 */
public class StatisticalTimeWindowFunction implements AllWindowFunction<networkData, networkData, TimeWindow> {

    private Logger logger = LoggerFactory.getLogger(StatisticalTimeWindowFunction.class);

    private HashMap<String, Integer> timeDstMap;
    private HashMap<String, Integer> timeServiceMap;


    public StatisticalTimeWindowFunction(HashMap<String, Integer> timeDstMap, HashMap<String, Integer> timeServiceMap) {
        this.timeDstMap = timeDstMap;
        this.timeServiceMap = timeServiceMap;
    }

    @Override
    public void apply(TimeWindow window, Iterable<networkData> values, Collector<networkData> out) throws Exception {
        int totalDstNumber = 0, totalServiceNumber = 0;
        for (String dst : timeDstMap.keySet()) {
            totalDstNumber += timeDstMap.get(dst);
        }
        for (String service : timeServiceMap.keySet()) {
            totalDstNumber += timeServiceMap.get(service);
        }
        for (networkData value : values) {
            if (!(value.getInternetLayer() instanceof IP)) {
                // 如果不是IP协议 不用计算，直接输入值
                value.setTimeWindowSameDstCount(0);
                value.setTimeWindowSameDstPercent(0);
                value.setTimeWindowDiffDstCount(totalDstNumber);
                value.setTimeWindowDiffDstPercent(1);
                value.setTimeWindowSameServiceCount(0);
                out.collect(value);
                continue;
            }

            IP ip = (IP) value.getInternetLayer();
            if (timeDstMap.containsKey(ip.getDstIP().toString())) {
                // 一般一定在相同地址中含有
                int count = timeDstMap.get(ip.getDstIP().toString());
                value.setTimeWindowSameDstCount(count);
                value.setTimeWindowSameDstPercent(((float) count) / ((float) totalDstNumber));
                value.setTimeWindowDiffDstCount((totalDstNumber - count));
                value.setTimeWindowDiffDstPercent(((float) (totalDstNumber - count)) / ((float) totalDstNumber));

            } else {
                // 如果发生了这种情况，说明计算出错了
                logger.error("目的地址map中不含有当前IP地址，请检查代码错误！！");
                value.setTimeWindowSameDstCount(0);
                value.setTimeWindowSameDstPercent(0);
                value.setTimeWindowDiffDstCount(totalDstNumber);
                value.setTimeWindowDiffDstPercent(1);
                value.setTimeWindowSameServiceCount(0);
            }
            if (value.getTransportLayer() == null) {
                // 不属于TCP和UDP协议，直接赋予值
                value.setTimeWindowSameServiceCount(0);
                out.collect(value);
                continue;
            }
            if (value.getTransportLayer() instanceof TCP) {
                // 传输层是TCP协议
                TCP tcp = (TCP) value.getTransportLayer();
                if (tcp.getApplicationProtocol().toString().compareTo("unknown") == 0) {
                    setTimeWindowSameServiceNumberTCP(value,tcp);
                } else {
                    setTimeWindowSameServiceNumberTCP(value,tcp);
                }
            } else {
                // 传输层是UDP协议
                UDP udp = (UDP) value.getTransportLayer();
                if (udp.getApplicationProtocol().toString().compareTo("unknown") == 0) {
                    setTimeWindowSameServiceNumberUDP(value,udp);
                } else {
                    setTimeWindowSameServiceNumberUDP(value,udp);
                }
            }
            out.collect(value);
        }

        timeDstMap.clear();
        timeServiceMap.clear();
    }

    public void setTimeWindowSameServiceNumberTCP(networkData value,TCP tcp){
        if(timeServiceMap.containsKey(tcp.getDstPort().toString())){
            value.setCountWindowSameServiceCount(timeServiceMap.get(tcp.getDstPort().toString()));
        }else{
            value.setCountWindowSameServiceCount(0);
        }
    }

    public void setTimeWindowSameServiceNumberUDP(networkData value,UDP udp){
        if(timeServiceMap.containsKey(udp.getDstPort().toString())){
            value.setCountWindowSameServiceCount(timeServiceMap.get(udp.getDstPort().toString()));
        }else{
            value.setCountWindowSameServiceCount(0);
        }
    }
}